"""
Variational Autoencoder (VAE) Module

实现结构化VAE，嵌入可微分博弈求解器
"""

from .encoder import Encoder, GaussianEncoder
from .decoder import Decoder, GameDecoder, DroneMCPDecoder, MyNormLayer
from .vae_model import MCP_VAE, compute_elbo_loss, DroneVAE
from .training import (
    VAETrainer,
    train_epoch,
    validate,
    log_solver_stats_to_file,
    sample_and_visualize_latent_space,
)

__all__ = [
    "Encoder",
    "GaussianEncoder",
    "Decoder",
    "GameDecoder",
    "DroneMCPDecoder",
    "MyNormLayer",
    "MCP_VAE",
    "DroneVAE",
    "VAETrainer",
    "compute_elbo_loss",
    "train_epoch",
    "validate",
    "log_solver_stats_to_file",
    "sample_and_visualize_latent_space",
]
